package tools

import (
	"context"
	"encoding/json"
	"log"

	"github.com/mark3labs/mcp-go/mcp"
)

const (
	SQLQueryToolName = "sql_query"
)

var SQLQueryTool = mcp.NewTool(
	SQLQueryToolName,
	mcp.WithDescription("Execute sql query statement, return a json data result.\n执行SQL查询语句并返回JSON结构的结果."),
	mcp.WithString(
		"sql",
		mcp.Description("Please provide a specific sql query statement conforming to the PostgreSQL syntax.\n请提供符合PostgreSQL语法规范的SQL查询语句."),
		mcp.Required(),
	),
)

func ExecuteSQLQueryHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	// 获取SQL语句
	sql := request.Params.Arguments["sql"].(string)
	log.Printf("REQUEST: %s", request.Method+", sql="+sql)
	// 查询数据库
	rows, _ := db.Query(ctx, sql)
	defer rows.Close()
	// 判断rows是否为空
	if rows == nil {
		return mcp.NewToolResultText("[]"), nil
	}
	// 获取列名
	cols := rows.FieldDescriptions()
	colNames := make([]string, len(cols))
	for i, col := range cols {
		colNames[i] = string(col.Name)
	}
	// 定义返回结构
	var datas []map[string]interface{}
	// 遍历查询结果
	for rows.Next() {
		// 创建一个切片存储每行的值（使用 interface{} 类型）
		values := make([]interface{}, len(cols))
		valuePtrs := make([]interface{}, len(cols))

		// 为每一列分配指针
		for i := range values {
			valuePtrs[i] = &values[i]
		}

		// 扫描当前行的值
		if err := rows.Scan(valuePtrs...); err != nil {
			return nil, err
		}

		// 将列名和对应的值存入 map
		rowMap := make(map[string]interface{})
		for i, col := range cols {
			val := values[i]

			// 处理数据库中的 NULL 值
			if b, ok := val.([]byte); ok {
				rowMap[col.Name] = string(b) // 将字节数组转换为字符串
			} else {
				rowMap[col.Name] = val
			}
		}

		// 将当前行的 map 添加到结果中
		datas = append(datas, rowMap)
	}

	result, _ := json.MarshalIndent(datas, "", "  ")

	return mcp.NewToolResultText(string(result)), nil
}
